from io import StringIO
import requests
import time

from loguru import logger

from models.table_defs import PaperReference

getYMDHmsRand = lambda  : int(round(time.time() * 1000))

class QryByUnstructCitation:
  usr = '934096384@qq.com'
  url0 = 'https://doi.crossref.org/servlet/query'

  @classmethod
  def str_to_one_line(cls, s):
    # 将多行字符串转为一行
    return ''.join(s.split('\n'))

  @classmethod
  def build_qdata_xml_str(cls, unstructured_citations):
    from lxml import etree
    doi_batch_id = getYMDHmsRand()
    query_batch_xml_str = """<query_batch xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" version="2.0" xmlns="http://www.crossref.org/qschema/2.0" xsi:schemaLocation="http://www.crossref.org/qschema/2.0 http://www.crossref.org/qschema/crossref_query_input2.0.xsd" > """
    xml_data = f'''<?xml version="1.0" ?>
    <query_batch><head><email_address>934096384@qq.com</email_address><doi_batch_id>{doi_batch_id}</doi_batch_id></head><body/></query_batch>
    '''
    f = StringIO(xml_data) # 如果把 query_batch_xml_str 的内容放到 xml_data 里，会导致后续 的 xpath 失灵
    tree = etree.parse(f)
    root = tree.getroot()
    bodyEl = root.xpath('//body')[0]
    for i, unstructured_citation in enumerate(unstructured_citations):
      queryEl = etree.SubElement(bodyEl, "query")
      queryEl.set("key", f"q{i + 1}")
      queryEl.set("enable-multiple-hits", "false")  # enable-multiple-hits="false"
      citation_el = etree.SubElement(queryEl, "unstructured_citation")
      citation_el.text = unstructured_citation

    xmlStr = etree.tostring(tree, encoding="utf-8", xml_declaration=True, pretty_print=False).decode("utf-8")
    xmlStr = xmlStr.replace('<query_batch>', query_batch_xml_str)
    return xmlStr

  @classmethod
  def query(cls, unstructured_citation):
    unstruct_citation = cls.str_to_one_line(unstructured_citation)
    params = {'usr': cls.usr, 'qdata': unstruct_citation, 'format': 'unixref'}
    response = requests.get(cls.url0, params=params)
    if response.status_code == 200:
      return response.text

  @classmethod
  def xml_to_PaperReference(cls, elem):
    paper_ref = PaperReference()
    paper_ref.title = elem.findtext('.//journal_article//title')

    # 提取作者
    authors = []
    for author in elem.findall('.//journal_article//person_name'):
      given_name = author.findtext('given_name')
      surname = author.findtext('surname')
      authors.append(f"{given_name} {surname}")
    paper_ref.authors = ' , '.join(authors)

    # 提取期刊名称
    journal_title = elem.findtext('.//full_title')
    paper_ref.journal = journal_title

    # 提取出版日期
    year = elem.findtext('.//year')
    month = elem.findtext('.//month')
    day = elem.findtext('.//day')
    paper_ref.publication_date = f"{year}-{month}-{day}"
    paper_ref.year = f"{year}"

    # 提取页码
    first_page = elem.findtext('.//first_page')
    last_page = elem.findtext('.//last_page')
    paper_ref.pages = f"{first_page}-{last_page}"

    # 提取 DOI 和资源
    paper_ref.doi = elem.findtext('.//doi')
    paper_ref.resource = elem.findtext('.//resource')

    return paper_ref

  @classmethod
  def query_4_PaperReference(cls,unstructured_citation):
    from lxml import etree

    res_xml = cls.query(unstructured_citation)
    if res_xml:
      # 将字符串编码为字节类型
      res_xml_bytes = res_xml.encode('utf-8')
      root = etree.fromstring(res_xml_bytes)
      obj_list = []
      for doi_rec in root.xpath('doi_record'):
        journal_elem = doi_rec.xpath('crossref/journal')[0]
        if journal_elem is not None:
          paper_ref=cls.xml_to_PaperReference(journal_elem)
          obj_list.append(paper_ref)
      return obj_list
    else:
      return []

# test code
def test_build_qdata_xml_str():
  unstructured_citations = [
    "Barsoum RS. Chronic kidney disease in the developing world. N Engl J Med 2006; 354: 997-9.",
    "Xu Y, Wang L, He J, et al. Prevalence and control of diabetes in Chinese adults. JAMA 2013; 310: 948-59.",
    "Yang W, Lu J, Weng J, et al. Prevalence of diabetes among men and women in China. N Engl J Med 2010; 362: 1090-101",
    "Pan XR, Yang WY, Li GW, Liu J. Prevalence of diabetes and its risk factors in China, 1994. Diabetes Care 1997; 20: 1664-9.",
    ]
  qdata_xml_str = QryByUnstructCitation.build_qdata_xml_str(unstructured_citations)
  print(qdata_xml_str)

def test_query():
  unstructured_citations = [
    "Barsoum RS. Chronic kidney disease in the developing world. N Engl J Med 2006; 354: 997-9.",
    "Xu Y, Wang L, He J, et al. Prevalence and control of diabetes in Chinese adults. JAMA 2013; 310: 948-59.",
    "Yang W, Lu J, Weng J, et al. Prevalence of diabetes among men and women in China. N Engl J Med 2010; 362: 1090-101",
    "Pan XR, Yang WY, Li GW, Liu J. Prevalence of diabetes and its risk factors in China, 1994. Diabetes Care 1997; 20: 1664-9.",
    ]
  qdata_xml_str = QryByUnstructCitation.build_qdata_xml_str(unstructured_citations)

  print(qdata_xml_str)
  print("\n\n")
  res_xml = QryByUnstructCitation.query(qdata_xml_str)
  logger.info(res_xml)

def test_query_4_PaperReference():
  unstructured_citations = [
    "Barsoum RS. Chronic kidney disease in the developing world. N Engl J Med 2006; 354: 997-9.",
    "Xu Y, Wang L, He J, et al. Prevalence and control of diabetes in Chinese adults. JAMA 2013; 310: 948-59.",
    "Yang W, Lu J, Weng J, et al. Prevalence of diabetes among men and women in China. N Engl J Med 2010; 362: 1090-101",
    "Pan XR, Yang WY, Li GW, Liu J. Prevalence of diabetes and its risk factors in China, 1994. Diabetes Care 1997; 20: 1664-9.",
    ]
  qdata_xml_str = QryByUnstructCitation.build_qdata_xml_str(unstructured_citations)
  QryByUnstructCitation.query_4_PaperReference(qdata_xml_str)

if __name__ == '__main__':
  # test_build_qdata_xml_str()
  # test_query()
  test_query_4_PaperReference()

"""
todo : 分批查询引用文献，每次最多5篇
"""
